# -*- coding: utf-8 -*-
# created on 2016/7/15
from mathsolver.functions.root.jiefangchenzu import JieFangChenZu
from mathsolver.functions.base import *
from collections import defaultdict
from mathsolver.functions.fangcheng import FangChenBuDengShi001, FangChenBuDengShi002
from mathsolver.functions.shuyushi import ShuYuShiFenShi001
from sympy import sympify, S, Set
from mathsolver.functions.mathematica.mathematicaSolve import MathematicaSolve
from mathsolver.functions.hanshu.zhiyu import HanShuZhiYu


# 计算集合
class JiHeCompute(BaseFunction):
    def final(self, new_set):
        self.output.append(new_set)
        self.steps.append(["", new_set.printing()])
        return self

    def solver(self, *args):
        if isinstance(args[0].value, list):
            self.output.append(BaseInnerSet({"name": "", "var": "", "value": FiniteSet(*args[0].value), "domain": "R"}))
            return self
        name = args[0].value["name"]
        var = args[0].value["var"]
        conditions = args[0].value["conditions"]
        is_or = args[0].value["or"] if "or" in args[0].value else 0
        # --> load data
        data = defaultdict(list)
        conditions = [cond for cond in conditions if not ("belong" in cond and cond["belong"][-1] == "R")]
        for cond in conditions:
            t = list(cond.keys())[0]
            data[t].append(cond[t])

        # M={-1,0,1},N={\\sqrt{1-{x}^{2}}|x∈M}
        # M={x|-1<x<2},N={y|y=\\frac{1}{2}{x}^{2}-1,x∈M},则M∩N=.
        flag = "R"
        if len(data["belong"]) <= 1:
            previous_condition = None
            for belong in data["belong"]:
                if belong[-1] in ["Z", "z"]:
                    flag = "Z"
                elif belong[-1] in ["N", "NN"]:
                    flag = "N"
                elif belong[-1] in ["N^{*}", "N_{+}"]:
                    flag = "N^*"
                elif belong[-1] == "R":
                    continue
                elif isinstance(belong[-1], list):
                    previous_condition = BaseInnerSet(
                        {'var': belong[0], 'domain': "R", "name": '', 'value': BaseInter(belong[-1]).interval})
                else:
                    previous_condition = self.search(belong[-1])
                    previous_condition.value["var"] = belong[0]
        else:
            single_belong = [belong for belong in data["belong"] if sympify(belong[0]).is_Symbol and belong[-1] != "R"]
            if len(single_belong) == 1:
                temp_var = sympify(single_belong[0][0])
                value = self.search(single_belong[0][-1]).value["value"]
                assert isinstance(value, list)
            else:
                temp_var = [sympify(belong[0]) for belong in single_belong]
                values = [self.search(belong[-1]).value["value"] for belong in single_belong]
                value = list(itertools.product(*values))

            poly_belongs = [belong for belong in data["belong"] if not sympify(belong[0]).is_Symbol]
            ignore_index = set()
            for poly_belong in poly_belongs:
                poly = sympify(poly_belong[0])
                target_value = self.search(poly_belong[-1]).value["value"]
                for i, v in enumerate(value):
                    symbol_value = dict(zip(temp_var, v)) if isinstance(temp_var, list) else {temp_var: v}
                    if poly.subs(symbol_value) not in target_value:
                        ignore_index.add(i)

            value = [v for i, v in enumerate(value) if i not in ignore_index]
            previous_condition = BaseInnerSet({'var': temp_var, 'domain': "R", "name": '', 'value': value})

        # --> start compute
        condition = conditions[0]
        if isinstance(var, list) and "eq" in condition and len(conditions) == 1:  # 解析几何 M={(x,y)|y=x^{2}}
            return self.final(
                BaseInnerSet({"name": name, "var": var, "value": BaseEq(condition["eq"]), "domain": flag}))
        elif isinstance(var, list) and previous_condition:
            return self.final(BaseInnerSet(
                {"name": name, "var": var, "value": FiniteSet(*previous_condition.value["value"]), "domain": flag}))
        elif previous_condition and isinstance(previous_condition.value["var"], list) and isinstance(
                previous_condition.value["value"], list) and "eq" not in condition and "ineq" not in condition:
            symbol_values = [dict(zip(previous_condition.value["var"], v)) for v in previous_condition.value["value"]]
            return self.final(BaseInnerSet(
                {"name": name, "var": "", "value": FiniteSet(*[sympify(var).subs(sym_v) for sym_v in symbol_values]),
                 "domain": flag}))

        elif "set" in condition:  # N={0,1,2}
            return self.final(BaseInnerSet(
                {"name": name, "var": var, "value": [sympify(s) for s in condition["set"]], "domain": flag}))

        stepsolver = None
        answer = None

        # 求定义域 B={x|y=lg(1-x)}
        if "eq" in condition and "y" == condition["eq"][0] and "x" == var and len(conditions) == 1:
            data["eq"].remove(condition["eq"])
            try:
                stepsolver = ShuYuShiFenShi001().solver(BasePoly(condition["eq"][1]))
                for item in stepsolver.output:
                    if isinstance(item, BaseIneq):
                        data["ineq"].append(item.value)
                    elif isinstance(item, BaseIneqs):
                        for ineq in item.ineqs:
                            if isinstance(ineq, BaseEq):
                                data["eq"].append(ineq.value)
                            else:
                                data["ineq"].append(ineq.value)
                    self.label.add("求函数定义域")
            except Exception:
                pass

        if len(data["eq"]) == 1 and "eq" in condition and var == condition["eq"][0] and "n" in condition["eq"][1]:
            f = sympify(condition["eq"][1])
            answer = FiniteSet(*[f.subs({Symbol("n"): i}) for i in range(20)])
        elif previous_condition and isinstance(previous_condition.value['value'], list) and len(
                data["eq"]) == 1 and var == condition["eq"][0]:
            f = sympify(data['eq'][0][1])
            temp_var = previous_condition.value['var']
            answer = []
            for v in previous_condition.value['value']:
                symbol_value = dict(zip(temp_var, v)) if isinstance(temp_var, list) else {temp_var: v}
                answer.append(f.subs(symbol_value))

            answer = FiniteSet(*answer)

        elif "eq" in condition and var == condition["eq"][0] and len(
                sympify(condition["eq"][1]).free_symbols) > 0:  # 求值域 A={y|y=x^{2}-1}
            cond = None
            if len(data["ineq"]) == 1:
                cond = BaseIneq(data["ineq"][0])
            elif len(data["ineqs"]) > 1:
                cond = BaseIneqs(data["ineq"])
            elif previous_condition:
                cond = previous_condition.value["value"]

            eq = condition["eq"]
            if str(eq[0]) == "x" and "x" not in str(eq[1]):
                eq[0] = "y"

            if cond:
                stepsolver = HanShuZhiYu().solver(BaseEq(eq).eq2func(), cond)
            else:
                try:
                    stepsolver = HanShuZhiYu().solver(BaseEq(eq).eq2func())
                except Exception:
                    pass

            answer = stepsolver.output[0].value
            self.label.update(stepsolver.label)

        elif len(data["eq"]) == 1 and len(data["ineq"]) == 0:  # 解方程 A={x||x-6|≤2}
            eq = BaseEq(data["eq"][0])
            if var:
                stepsolver = MathematicaSolve().solver(eq, BaseVariable(var))
            else:
                stepsolver = MathematicaSolve().solver(eq)

            self.label.add("解方程")

        elif len(data["eq"]) > 1 and len(data["ineq"]) == 0:  # 解方程组
            eqs = BaseEqs(data["eq"])
            stepsolver = JieFangChenZu().solver(eqs)

        elif len(data["ineq"]) == 1 and len(data["eq"]) == 0:  # 解不等式 B={x|{x}^{2}-4x+3≤0}
            # ineq = data["ineq"][0]
            # if len(ineq) == 3:
            #     try:
            #         stepsolver = fangchen_budengshi_001().solver(BaseIneq(ineq), BaseVariable(var))
            #     except:
            #         answer = ineq
            # else:
            try:
                if re.search("[fgh]\\(.*\\)", str(data["ineq"][0]).replace("log", "t")):
                    from mathsolver.functions.process import ProcessExpression
                    ineq = ProcessExpression(known=self.known).solver(BaseValue(data["ineq"][0])).output[0]
                else:
                    ineq = BaseIneqs(data["ineq"])
                if var:
                    stepsolver = MathematicaSolve().solver(ineq, BaseVariable(var))
                else:
                    stepsolver = MathematicaSolve().solver(ineq)

                self.label.add("解不等式")
            except Exception:
                answer = BaseIneqs(data["ineq"][0])

        elif len(data["ineq"]) > 1 and len(data["eq"]) == 0:  # 解不等式组 N={x|\\frac{1}{4}<(\\frac{1}{2})^{x}<1}
            if is_or:
                answer = None
                for ineq in data["ineq"]:
                    if len(ineq) == 3:
                        ineq = BaseIneq(ineq)
                        stepsolver = FangChenBuDengShi001().solver(ineq, BaseVariable(var))
                    else:
                        stepsolver = FangChenBuDengShi002().solver(BaseIneqs(ineq), BaseVariable(var))
                    value = stepsolver.output[0].value.values()[0]
                    self.label.update(stepsolver.label)
                    if answer is None:
                        answer = value
                    else:
                        answer = answer.union(value)

            else:
                ineqs = BaseIneqs(data["ineq"])
                stepsolver = FangChenBuDengShi002().solver(ineqs, BaseVariable(var))

        elif len(data["ineq"]) > 0 and len(data["eq"]) > 0:
            ineqs = BaseIneqs(data["ineq"] + data["eq"])
            stepsolver = FangChenBuDengShi002().solver(ineqs)

        if not answer:
            if stepsolver:
                self.label.update(stepsolver.label)
                answer = stepsolver.output[0]
            else:
                answer = S.Reals

        if isinstance(answer, BaseSymbolValue) and flag in ["Z", "N", "N^*"]:
            key, value = list(answer.value.items())[0]
            ans = value.intersection(S.Integers)
            if flag == "N":
                ans = ans.intersection(Interval(0, S.Infinity))
            elif flag == "N^*":
                ans = ans.intersection(Interval(0, S.Infinity, True))
            answer = BaseSymbolValue({key: list(ans)})
            self.label.add("不等式的特殊整数解问题")

        if isinstance(answer, BaseSymbolValue):
            value = answer.value
            assert len(value) == 1  # and (str(value.keys()[0]) == var or var == ""), "not match"
            v = list(value.values())[0]
            if isinstance(v, list) and len(v) == 0:
                v = EmptySet
            var = list(value.keys())[0]
            return self.final(BaseInnerSet({"name": name, "var": var, "value": v, "domain": flag}))
        elif isinstance(answer, FiniteSet) and len(answer) > 10:
            self.output.append(BaseInnerSet({"name": name, "var": var, "value": answer, "domain": flag}))
            return self
        elif isinstance(answer, Set) or isinstance(answer, BaseIneqs) or isinstance(answer, BaseIneq):
            return self.final(BaseInnerSet({"name": name, "var": var, "value": answer, "domain": flag}))
            # return self
        else:
            raise ValueError("unexpected situation")


class JiHeSets(BaseFunction):
    def solver(self, *args):
        sets = args[0].value
        for tmp in sets:
            stepsolver = JiHeCompute().solver(BaseSet(tmp))
            self.output += stepsolver.output
            self.label.update(stepsolver.label)
            self.steps += stepsolver.steps
        return self


if __name__ == '__main__':
    pass
